// jpge.h - C++ class for JPEG compression.
// Public domain, Rich Geldreich <richgel99@gmail.com>
// Alex Evans: Added RGBA support, linear memory allocator.
#ifndef JPEG_ENCODER_H
#define JPEG_ENCODER_H

#include <stdint.h>

namespace jpge {
typedef char int8;
typedef unsigned char uint8;
typedef signed short int16;
typedef signed int int32;
typedef unsigned short uint16;
typedef unsigned int uint32;
typedef unsigned int uint;

// JPEG chroma subsampling factors. Y_ONLY (grayscale images) and H2V2 (color
// images) are the most common.
enum subsampling_t { Y_ONLY = 0, H1V1 = 1, H2V1 = 2, H2V2 = 3 };

// JPEG compression parameters structure.
struct params {
    inline params() : m_quality(85), m_subsampling(H2V2) {
    }

    inline bool check() const {
        if ((m_quality < 1) || (m_quality > 100)) {
            return false;
        }
        if ((uint)m_subsampling > (uint)H2V2) {
            return false;
        }
        return true;
    }

    // Quality: 1-100, higher is better. Typical values are around 50-95.
    int m_quality;

    // m_subsampling:
    // 0 = Y (grayscale) only
    // 1 = H1V1 subsampling (YCbCr 1x1x1, 3 blocks per MCU)
    // 2 = H2V1 subsampling (YCbCr 2x1x1, 4 blocks per MCU)
    // 3 = H2V2 subsampling (YCbCr 4x1x1, 6 blocks per MCU-- very common)
    subsampling_t m_subsampling;
};

// Output stream abstract class - used by the jpeg_encoder class to write to the
// output stream. put_buf() is generally called with len==JPGE_OUT_BUF_SIZE
// bytes, but for headers it'll be called with smaller amounts.
class output_stream {
   public:
    virtual ~output_stream(){};
    virtual bool put_buf(const void *Pbuf, int len) = 0;
    virtual uint get_size() const                   = 0;
};

// Lower level jpeg_encoder class - useful if more control is needed than the
// above helper functions.
class jpeg_encoder {
   public:
    jpeg_encoder();
    ~jpeg_encoder();

    // Initializes the compressor.
    // pStream: The stream object to use for writing compressed data.
    // params - Compression parameters structure, defined above.
    // width, height  - Image dimensions.
    // channels - May be 1, or 3. 1 indicates grayscale, 3 indicates RGB source
    // data. Returns false on out of memory or if a stream write fails.
    bool init(output_stream *pStream, int width, int height, int src_channels,
              const params &comp_params = params());
    bool reinit(int quality);

    // Call this method with each source scanline.
    // width * src_channels bytes per scanline is expected (RGB or Y format).
    // You must call with NULL after all scanlines are processed to finish
    // compression. Returns false on out of memory or if a stream write fails.
    bool process_scanline(const void *pScanline);
    bool process_scanline565(const void *pScanline);
    void process_mcu_row();

    // Deinitializes the compressor, freeing any allocated memory. May be called
    // at any time.
    void deinit();

   private:
    jpeg_encoder(const jpeg_encoder &);
    jpeg_encoder &operator=(const jpeg_encoder &);

    typedef int32 sample_array_t;
    enum { JPGE_OUT_BUF_SIZE = 1500, JPGE_OUT_BUF_COUNT = 4 };

    output_stream *m_pStream;
    params m_params;
    uint8 m_num_components;
    uint8 m_comp_h_samp[3], m_comp_v_samp[3];
    int m_image_x, m_image_y, m_image_bpp, m_image_bpl;
    int m_image_x_mcu, m_image_y_mcu;
    int m_image_bpl_xlt, m_image_bpl_mcu;
    int m_mcus_per_row;
    int m_mcu_x, m_mcu_y;
    int16 *m_mcu_lines[16];
    uint8 m_mcu_y_ofs;
    sample_array_t m_sample_array[64];
    int16 m_coefficient_array[64];

    int32 m_last_dc_val[3];
    uint8 m_out_buf_array[JPGE_OUT_BUF_COUNT][JPGE_OUT_BUF_SIZE];
    uint8 m_out_buf_index;
    uint8 *m_pOut_buf;
    uint m_out_buf_left;
    uint32 m_bit_buffer;
    uint m_bits_in;
    uint8 m_pass_num;
    bool m_all_stream_writes_succeeded;

    bool jpg_open(int p_x_res, int p_y_res, int src_channels);

    void flush_output_buffer();
    void put_bits(uint bits, uint len);

    void emit_byte(uint i);
    void emit_word(uint i);
    void emit_marker(uint8 marker);

    void emit_jfif_app0();
    void emit_dqt();
    void emit_sof();
    void emit_dht(uint8 *bits, uint8 *val, int index, bool ac_flag);
    void emit_dhts();
    void emit_sos();

    void compute_quant_table(uint8 *pDst, int32 *pDst32, const uint8 *pSrc);
    void load_quantized_coefficients(int component_num);

    void load_block_8_8_grey(int x);
    void load_block_8_8(int x, int y, int c);
    void load_block_16_8(int x, int c);
    void load_block_16_8_8(int x, int c);

    void code_coefficients_pass_two(int component_num);
    void code_block(int component_num);

    // void process_mcu_row();
    bool process_end_of_image();
    void load_mcu(const void *src);
    void load_mcu565(const void *src);
    void clear();
    void init();
};

}  // namespace jpge

#endif  // JPEG_ENCODER
